package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;


import nlp.langmodel.LanguageModel;
import nlp.util.Counter;
public class TrigramLanguageModel implements LanguageModel {
	private static String START = "<S>";
	private static String STOP = "</S>";
	private static double INCREMENT_FACTOR = 1.0;
	// unigram should be set to normal counter instead of GoodTuringCounter
	private GoodTuringCounter<String> unigramCounter = new GoodTuringCounter<String>();
	private GoodTuringCounter<Bigram<String, String>> bigramCounter = new GoodTuringCounter<Bigram<String,String>>();
	private Counter<String> specialUnigramCounter = new Counter<String>();
	private GoodTuringCounter<Trigram<String, String, String>> trigramCounter = new GoodTuringCounter<Trigram<String,String,String>>();
	private Counter<Bigram<String, String>> specialBigramCounter = new Counter<Bigram<String,String>>();
	private Collection<List<String>> validationCollection;
	Trigram<Double, Double, Double> lambdas = null;
	private double unigramTotal;
	private double bigramTotal;
	private double trigramTotal;
	
	public TrigramLanguageModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.validationCollection = validationCollection;
		// unigram generation
		for(List<String> sentence: trainingCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(STOP);
			for (String word : stoppedSentence) {
				unigramCounter.incrementCount(word, INCREMENT_FACTOR);
			}
		}
		
		// bigram generation
		for (List<String> sentence: trainingCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(0, START);
			stoppedSentence.add(STOP);
			
			Iterator<String> iter = stoppedSentence.iterator();
			String first = iter.next();
			String second = iter.next();
			while (iter.hasNext()) {
				Bigram<String, String> bigram = new Bigram<String, String>(first, second);
				bigramCounter.incrementCount(bigram, INCREMENT_FACTOR);
				first = second;
				if (iter.hasNext()) {
					second = iter.next();
				}
			}
		}
		unigramCounter.smooth();
		bigramCounter.smooth();
		
		for (Bigram<String, String> bigram : bigramCounter.keySet()) {
			specialUnigramCounter.incrementCount(bigram.first, INCREMENT_FACTOR * bigramCounter.getCount(bigram));
		}

		// trigram generation
		for (List<String> sentence: trainingCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(0, START);
			stoppedSentence.add(0, START);
			stoppedSentence.add(STOP);
			
			Iterator<String> iter = stoppedSentence.iterator();
			String first = iter.next();
			String second = iter.next();
			String third = iter.next();
			
			while (iter.hasNext()) {
				Trigram<String, String, String> trigram = new Trigram<String, String, String>(first, second, third);
				trigramCounter.incrementCount(trigram, INCREMENT_FACTOR);
				first = second;
				second = third;
				if (iter.hasNext()) {
					third = iter.next();
				}
			}
		}
		trigramCounter.smooth();
		
		for (Trigram<String, String, String> trigram : trigramCounter.keySet()) {
			Bigram<String, String> bigram = new Bigram<String, String>(trigram.first, trigram.second);
			specialBigramCounter.incrementCount(bigram, INCREMENT_FACTOR * trigramCounter.getCount(trigram));
		}
		
		// smooth counter
		
		unigramTotal = unigramCounter.totalCount();
		bigramTotal = bigramCounter.totalCount();
		trigramTotal = trigramCounter.totalCount();
		System.out.println("uniCount = " + unigramTotal);
		System.out.println("bigramCount = " + bigramTotal);
		System.out.println("trigramCount = " + trigramTotal);
		
		// train validation set
//		lambdas = resolveLambdas();
//		lambdaPair = new Pair<Double, Double>(0.0, 1.0);
//		System.out.println("lambdaPair is " + lambdas);
		System.out.println("LanguageModelTester.calculatePerplexity() is " + LanguageModelTester.calculatePerplexity(this, validationCollection));
	}
	
	
	private Trigram<Double, Double, Double> resolveLambdas() {
		double gap = .1;
		double lambda1 = 1.0;
		double lambda2 = 0.0;
		double lambda3 = 0.0;
		int times = (int)(1.0 / gap);
		double perplexity = Double.POSITIVE_INFINITY;
		for (int i = 1; i < times; i++) {
			for (int j = 1; j < times - i; j++) {
				double t1 = doMul(i, gap);
				double t2 = doMul(j, gap);
				double t3 = 1.0 - t1 - t2;
				System.out.println(t1 + ", " + t2 + ", " + t3);
				double temp = calculatePerplexity(this.validationCollection, t1, t2, t3);
				if (perplexity > temp) {
					perplexity = temp;
					lambda1 = t1;
					lambda2 = t2;
					lambda3 = t3;
				}
//				System.out.println("temp perplexity of rount " + (i * times + j) + " is " + temp);
			}
		}
		System.out.println("\nlowest perplexity " + perplexity);

		return new Trigram<Double, Double, Double>(lambda1, lambda2, lambda3);
	}
	
	private double calculatePerplexity(
			Collection<List<String>> sentenceCollection, double l1, double l2,
			double l3) {
		double logProbability = 0.0;
		double numSymbols = 0.0;
		for (List<String> sentence : sentenceCollection) {
			logProbability += getSentenceLogProbability(sentence,
					l1, l2, l3);
			numSymbols += sentence.size();
		}
		double avgLogProbability = logProbability / numSymbols;
		System.out.println("avgLogProbability is " + avgLogProbability);
		double perplexity = Math.pow(.5, avgLogProbability);
		return perplexity;
	}

	private double getSentenceLogProbability(List<String> sentence,
			double l1, double l2, double l3) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		double logProbability = 0.0;
		stoppedSentence.add(0, START);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		for (int i = 0; i < stoppedSentence.size(); i++) {
			logProbability += getInterpolatedWordLogProbability(stoppedSentence, i, l1, l2, l3);
		}
		System.out.println("sentence p is " + logProbability);
		return logProbability;
	}
	
	private double getInterpolatedWordLogProbability(List<String> sentence, int index, double l1, double l2, double l3) {
		if (index <= 1) {
			return 0.0;
		}
		String third = sentence.get(index);
		String second = sentence.get(index - 1);
		String first = sentence.get(index - 2);
		Trigram<String, String, String> trigram = new Trigram<String, String, String>(first, second, third);
		
		double pTrigram = 0.0;
		double specialBigram = specialBigramCounter.getCount(new Bigram<String, String>(trigram.first, trigram.second));
		if (!trigramCounter.keySet().contains(trigram)) {
			pTrigram = trigramCounter.getZeroCount() / trigramTotal;
		} else {
			pTrigram = trigramCounter.getCount(trigram) / specialBigram;
		}
		
		Bigram<String, String> bigram = new Bigram<String, String>(second, third);
		double pBigram = 0.0;
		double specialUnigram = specialUnigramCounter.getCount(first);
		if (!bigramCounter.keySet().contains(bigram)) {
			pBigram = bigramCounter.getZeroCount() / bigramTotal;
		} else {
			pBigram = bigramCounter.getCount(bigram) / specialUnigram;
		}
		
		String unigram = third;
		double pUnigram = 0.0;
		if (!unigramCounter.keySet().contains(unigram)) {
			pUnigram = unigramCounter.getZeroCount() / unigramTotal;
		} else {
			pUnigram = unigramCounter.getCount(unigram) / unigramTotal;
		}
		
		double logProbability1 = doMul(l1, pUnigram);
		double logProbability2 = doMul(l2, pBigram);
		double logProbability3 = doMul(l3, pTrigram);
		double temp = logProbability1 + logProbability2 + logProbability3;
//		System.out.println("p1 is " + logProbability1 + ", p2 is " + logProbability2 + " and p3 is " + logProbability3);

		return  Math.log(temp) / Math.log(2.0);
	}
	
	private double getWordLogProbability(List<String> sentence, int index) {
		if (index <= 1) {
			return 0.0;
		}
		String third = sentence.get(index);
		String second = sentence.get(index - 1);
		String first = sentence.get(index - 2);
		Trigram<String, String, String> trigram = new Trigram<String, String, String>(first, second, third);
		
		double pTrigram = 0.0;
		double specialBigram = specialBigramCounter.getCount(new Bigram<String, String>(trigram.first, trigram.second));
		if (!trigramCounter.keySet().contains(trigram)) {
			pTrigram = trigramCounter.getZeroCount() / trigramTotal;
		} else {
			pTrigram = trigramCounter.getCount(trigram) / specialBigram;
		}
		
		return Math.log(pTrigram) / Math.log(2.0);
	}

	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		double logProbability = 0.0;
		stoppedSentence.add(0, START);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		for (int i = 0; i < stoppedSentence.size(); i++) {
			logProbability += getWordLogProbability(stoppedSentence, i);
		}
//		System.out.println("sentence p is " + logProbability);
		return Math.pow(2.0, logProbability);
	}

	private String generateWord(String first, String second) {
		double sample = Math.random();
		double sum = 0.0;
		Bigram<String, String> special = new Bigram(first, second);
		double specialBigram = specialBigramCounter.getCount(special);
		for (Trigram<String, String, String> trigram : trigramCounter.keySet()) {
			Bigram<String, String> bigram = new Bigram<String, String>(trigram.first, trigram.second);
			if (bigram.equals(special)) {
				sum += trigramCounter.getCount(trigram) / specialBigram;
			}
			if (sum > sample) {
				return trigram.third;
			}
		}		
		
		return STOP;
	}
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		sentence.add(START);
		sentence.add(START);
		String first = START;
		String second = generateWord(START, START);
		String next = second;
		while (!next.equals(STOP) && sentence.size() < 20) {
//			System.out.println(word);
			sentence.add(next);
			next = generateWord(first, second);
			first = second;
			second = next;
		}
		sentence.add(STOP);
		System.out.println("sentence size is " + sentence.size());
		return sentence;
	}
	
	static double doMul(double op1, double op2) {
		if (op1 == 0 || op2 == 0) {
			return 0.0;
		}
		return Math.pow(Math.E, Math.log(op1) + Math.log(op2));	
	}
	

	public static void main(String[] args) {
		System.out.println(TrigramLanguageModel.doMul(0.01, 0.1));
	}
}
